# -*- coding: utf-8 -*-
from __future__ import print_function

import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.optim as optim

from ddl_platform.ddlib import job
from datasets import DatasetHDF5
import os
#import models
import torchvision.models as models

class ImageNetVGGJob(job.Job):
    def build_dataset(self):
        config = self.config()['dataset']
        data_dir = config['data_dir']

        ## Data loading code
        traindir = os.path.join(data_dir, 'train')
        testdir = os.path.join(data_dir, 'val')

        normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                         std=[0.229, 0.224, 0.225])

        image_size = 224
        trainset = torchvision.datasets.ImageFolder(traindir, transforms.Compose([
        #hdf5fn = os.path.join(data_dir, 'imagenet-shuffled.hdf5')
        #self._input_shape = (config['batch_size'], 3, image_size, image_size)
        #self._output_shape = (config['batch_size'], 1000)
        #trainset = DatasetHDF5(hdf5fn, 'train', transforms.Compose([
        #    transforms.ToPILImage(),
            transforms.RandomResizedCrop(image_size),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize,
            ]))
        self.trainset = trainset

        #testset = DatasetHDF5(hdf5fn, 'val', transforms.Compose([
        #        transforms.ToPILImage(),
        testset = torchvision.datasets.ImageFolder(testdir, transforms.Compose([
                transforms.Resize(256),
                transforms.CenterCrop(224),
                transforms.ToTensor(),
                normalize,
            ]))

        return trainset, testset


    def build_model(self):
        #return models.VGG("VGG16", input_size=224, num_classes=1000)
        return models.vgg16()
        #return models.VGG("VGG16", input_size=224, num_classes=1000)

    def build_optimizer(self):
        model = self.model()
        config = self.config()['optimizer']
        optimizer = optim.SGD(model.parameters(), 
                lr=config['lr'],
                momentum=config['momentum'])
        return optimizer

    def build_criterion(self):
        return nn.CrossEntropyLoss()

    def cal_eval_performance(self, batch_outputs, batch_inputs):
        return super().cal_eval_performance(batch_outputs, batch_inputs)


#def train():
#    conf_yaml = 'mnist.yaml'
#    job = MnistJob(conf_yaml)
#    t = trainer.Trainer(job)
#    t.fit()
#
#if __name__ == '__main__':
#    train()
